# pip install zhipuai 请先在终端进行安装
import os

from zhipuai import ZhipuAI
import re
import subprocess

client = ZhipuAI(api_key=os.environ["ZHIPU_API_KEY"])

#glm-4大模型现在通过提示，是无法一次性写对有难度的代码的

# 提取Python代码
def extract_code_between_backticks(text):
    pattern = r"```python(.*?)```"
    match = re.search(pattern, text, re.DOTALL)
    if match:
        return match.group(1)
    else:
        return None

# 运行Python代码并检查结果
def run_python_script_and_check(script_path):
    try:
        result = subprocess.run(["python", script_path], capture_output=True, text=True, check=True)
        output = result.stdout
        return True,output
    except subprocess.CalledProcessError as e:
        error_output = e.stderr
        return False, error_output

prompt = """请写一段Pytorch代码。现在我想实现一个基于AlexNet的网络架构，\
这个网络架构有两部分输入，这两部分输入结构不一致，一种输入是128x128的图像\
另一种输入是256x256的图像，但在架构的中间两部分输入的结构会合并，\
这个架构一共只有7层，两种不同结构的数据在第四层合并，最后一层是线性层\
请你构建这个架构，并假设两个输入数据，验证这个网络能够顺利运行。\
无需进行训练，只要代码架构能够通就可以\
请仔细考虑两种不同结构的输入要如何融合，不要给我返回会报错的代码，谢谢。"""

def reflection(prompt):
    #反思次数
    reflectionTimes = 5

    codeNew = None

    coderMessages = [
        {
            "role": "system",
            "content": "你是一个资深的深度学习程序员"
        },
        {
            "role": "user",
            "content": prompt
        }
    ]

    coderResponse = client.chat.completions.create(
        model="glm-4",  # 填写需要调用的模型名称
        messages=coderMessages,
    )
    testPrompt = "请你对这段代码进行仔细的检查和回顾，观察这段代码能否顺利运行、是否会报错。深度学习代码很容易出现设置的超参数与数据实际情况不匹配的情况，请仔细检查，并让对方根据你的建议重写代码。"

    while reflectionTimes>0:
        codeNew = coderResponse.choices[0].message.content

        code = extract_code_between_backticks(codeNew)

        with open("temp_code.py", "w") as file:
            file.write(code)

        success,stdout = run_python_script_and_check("temp_code.py")

        if(success == True):
            print("===========",reflectionTimes,"代码运行成功", "\n")
            break
        else:
            print("===========",reflectionTimes,"代码运行失败", "\n")

        # print("===========", reflectionTimes, "产生代码",codeNew,"\n")

        checkMessages = [
            {
                "role": "system",
                "content": "你是一个资深的深度学习程序员，同时也是一个代码测试专家。你能够检查代码之中是否有错误导致代码不能运行，并就可能的错误给出建议"
            },
            {
                "role": "user",
                "content": prompt + "\n" + codeNew + "\n" + testPrompt
            },
        ]

        checkResponse = client.chat.completions.create(
            model="glm-4",  # 填写需要调用的模型名称
            messages=checkMessages
        )

        suggestion = checkResponse.choices[0].message.content
        #print("===========",reflectionTimes,"代码检查建议", suggestion, "\n")

        coderMessages.append({"role": "assistant","content": codeNew + "\n报错信息：" + stdout + "\n"})
        coderMessages.append({"role": "user", "content": suggestion + "\n请根据前面的建议完善代码，并返回修改后的完整代码给我\n"})
        coderResponse = client.chat.completions.create(
            model="glm-4",  # 填写需要调用的模型名称
            messages=coderMessages,
        )

        reflectionTimes-=1

    print("========最终代码============\n")
    print(codeNew)

if __name__ == '__main__':
    reflection(prompt)
